package com.ai.module.mj.controller;

import cn.hutool.core.comparator.CompareUtil;
import com.ai.module.chat.common.base.ThreadLocalContext;
import com.ai.module.mj.constant.enums.TaskStatus;
import com.ai.module.mj.domain.Task;
import com.ai.module.mj.domain.dto.TaskCondition;
import com.ai.module.mj.domain.req.RobotProcessTaskReq;
import com.ai.module.mj.domain.req.TaskConditionDTO;
import com.ai.module.mj.domain.vo.TaskStatusVo;
import com.ai.module.mj.loadbalancer.DiscordLoadBalancer;
import com.ai.module.mj.service.HighTaskStoreService;
import io.swagger.annotations.Api;
import io.swagger.annotations.ApiOperation;
import io.swagger.annotations.ApiParam;
import lombok.RequiredArgsConstructor;
import org.springframework.util.CollectionUtils;
import org.springframework.web.bind.annotation.*;

import java.util.*;
import java.util.stream.Collectors;

/**
 * @author: 何晋庆
 * @create: 2023-12-10 12:55
 * @description: 查询 任务队列数
 **/
@Api(tags = "任务查询")
@RestController
@RequestMapping("/member/task")
@RequiredArgsConstructor
public class TaskController {
	private final HighTaskStoreService taskStoreService;
	private final DiscordLoadBalancer discordLoadBalancer;

	@ApiOperation(value = "指定taskID获取任务")
	@GetMapping("/{id}/fetch")
	public Task fetch(@ApiParam(value = "任务ID") @PathVariable String id) {
		return this.taskStoreService.get(id);
	}

	@ApiOperation(value = "查询任务队列", hidden = true)
	@GetMapping("/queue")
	public List<Task> queue() {
		return this.discordLoadBalancer.getQueueTaskIds().stream()
				.map(this.taskStoreService::get).filter(Objects::nonNull)
				.sorted(Comparator.comparing(Task::getSubmitTime))
				.collect(Collectors.toList());
	}

	@ApiOperation(value = "查询所有任务", hidden = true)
	@GetMapping("/list")
	public List<Task> list() {
		return this.taskStoreService.list().stream()
				.sorted((t1, t2) -> CompareUtil.compare(t2.getSubmitTime(), t1.getSubmitTime()))
				.collect(Collectors.toList());
	}

	@ApiOperation(value = "根据ID列表查询任务", hidden = true)
	@PostMapping("/list-by-condition")
	public List<Task> listByIds(@RequestBody TaskConditionDTO conditionDTO) {
		if (conditionDTO.getIds() == null) {
			return Collections.emptyList();
		}
		return conditionDTO.getIds().stream().map(this.taskStoreService::get).filter(Objects::nonNull).collect(Collectors.toList());
	}

	@PostMapping("findProcessTask")
	public TaskStatusVo findProcessTask(@RequestBody RobotProcessTaskReq robotProcessTaskReq){
		TaskCondition condition = new TaskCondition()
				.setStatusSet(Set.of(TaskStatus.IN_PROGRESS,
						TaskStatus.SUBMITTED,
						TaskStatus.NOT_START))
				.setRobotId(robotProcessTaskReq.getRobotId())
				.setUserId(ThreadLocalContext.getUserId());
		List<Task> taskList = this.discordLoadBalancer.findRunningTask(condition).collect(Collectors.toList());
		if (! CollectionUtils.isEmpty(taskList)){
			TaskStatusVo taskStatusVo = new TaskStatusVo();
			List<TaskStatusVo.TaskProcess> taskProcesses = new ArrayList<>();
			taskList.forEach(task->{
				TaskStatusVo.TaskProcess taskProcess = new TaskStatusVo.TaskProcess();
				taskProcess.setStatus(task.getStatus());
				taskProcess.setProgress(task.getProgress());
				taskProcess.setTaskId(task.getId());
				taskProcesses.add(taskProcess);
			});
			taskStatusVo.setTaskProcessList(taskProcesses);
			return taskStatusVo;
		}
		return new TaskStatusVo();
	}
}
